import torch
import numpy as np
import audioread
import argparse
import time
import json
from torchaudio.transforms import Resample
import gushi_trainer

# --- 音频加载函数 (来自原项目) ---
def buf_to_float(x, n_bytes=2, dtype=np.float32):
    scale = 1.0 / float(1 << ((8 * n_bytes) - 1))
    fmt = f"<i{n_bytes}"
    return scale * np.frombuffer(x, fmt).astype(dtype)

def audioread_load(path, offset=0.0, duration=None, dtype=np.float32):
    y = []
    sr_native = 0
    with audioread.audio_open(path) as input_file:
        sr_native = input_file.samplerate
        n_channels = input_file.channels
        s_start = int(np.round(sr_native * offset)) * n_channels
        s_end = np.inf if duration is None else s_start + (int(np.round(sr_native * duration)) * n_channels)
        
        n = 0
        for frame in input_file:
            frame = buf_to_float(frame, dtype=dtype)
            n_prev = n
            n += len(frame)
            if n < s_start: continue
            if s_end < n_prev: break
            if s_end < n: frame = frame[:s_end - n_prev]
            if n_prev <= s_start <= n: frame = frame[s_start - n_prev:]
            y.append(frame)

    if not y: return np.empty(0, dtype=dtype), sr_native

    y = np.concatenate(y)
    if n_channels > 1:
        y = y.reshape((-1, n_channels)).T[0]
    return y, sr_native

# --- 主评估函数 ---
def perform_evaluation(audio_path: str, text: str, task: str):
    print(f"--- 古诗词'{task}'评测开始 ---")
    print(f"音频文件: {audio_path}")
    print(f"参考诗词: \"{text}\"")

    try:
        # 1. 加载和预处理音频
        signal_np, fs_orig = audioread_load(audio_path)
        if fs_orig == 0:
            print("错误: 无法加载音频或音频为空。")
            return

        signal_torch = torch.from_numpy(signal_np.astype(np.float32))

        if fs_orig != 16000:
            print(f"原始采样率: {fs_orig} Hz, 正在重采样至 16000 Hz...")
            resampler = Resample(orig_freq=fs_orig, new_freq=16000)
            signal_torch = resampler(signal_torch)
        
        if signal_torch.ndim == 1:
            audio_tensor = signal_torch.unsqueeze(0)
        else:
            audio_tensor = signal_torch

        # 2. 初始化评测器
        print("\n正在初始化评测模型...")
        start_load = time.time()
        trainer = gushi_trainer.get_gushi_trainer()
        print(f"模型加载完毕，耗时: {time.time() - start_load:.2f} 秒")

        # 3. 执行评测
        print("正在进行评测...")
        start_eval = time.time()
        result = trainer.evaluate(audio_tensor, text, task)
        print(f"评测完成，耗时: {time.time() - start_eval:.2f} 秒")

        # 4. 打印结果
        print("\n--- 评测结果 ---")
        # 使用json库美化输出
        print(json.dumps(result, indent=4, ensure_ascii=False))

    except FileNotFoundError:
        print(f"错误: 音频文件 '{audio_path}' 未找到。")
    except Exception as e:
        print(f"评估过程中发生未知错误: {e}")
        import traceback
        traceback.print_exc()

# --- 命令行接口 ---
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="中文古诗词朗读与背诵评测工具")
    parser.add_argument("audio_file", type=str, help="待评测的音频文件路径 (例如: 'gushi_audio.wav')")
    parser.add_argument("reference_text", type=str, help="参考的古诗词文本 (例如: '床前明月光')")
    parser.add_argument(
        "--task", 
        type=str, 
        default="reading", 
        choices=["reading", "memorization"], 
        help="选择评测任务: 'reading' (朗读水平评估) 或 'memorization' (背诵准确度检查)"
    )
    
    args = parser.parse_args()
    
    perform_evaluation(args.audio_file, args.reference_text, args.task) 